// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/kernels/row_encoder_internal.h>

#include <nebula/bits/bitmap_writer.h>
#include <turbo/log/logging.h>

#include <memory>

using nebula::internal::FirstTimeBitmapWriter;

namespace nebula::compute::internal {

// extract the null bitmap from the leading nullity bytes of encoded keys
    turbo::Status KeyEncoder::DecodeNulls(MemoryPool *pool, int32_t length, uint8_t **encoded_bytes,
                                          std::shared_ptr<Buffer> *null_bitmap,
                                          int32_t *null_count) {
        // first count nulls to determine if a null bitmap is necessary
        *null_count = 0;
        for (int32_t i = 0; i < length; ++i) {
            *null_count += (encoded_bytes[i][0] == kNullByte);
        }

        if (*null_count > 0) {
            TURBO_MOVE_OR_RAISE(*null_bitmap, allocate_bitmap(length, pool));
            uint8_t *validity = (*null_bitmap)->mutable_data();

            FirstTimeBitmapWriter writer(validity, 0, length);
            for (int32_t i = 0; i < length; ++i) {
                if (encoded_bytes[i][0] == kValidByte) {
                    writer.Set();
                } else {
                    writer.Clear();
                }
                writer.next();
                encoded_bytes[i] += 1;
            }
            writer.finish();
        } else {
            for (int32_t i = 0; i < length; ++i) {
                encoded_bytes[i] += 1;
            }
        }
        return turbo::OkStatus();
    }

    void BooleanKeyEncoder::AddLength(const ExecValue &, int64_t batch_length,
                                      int32_t *lengths) {
        for (int64_t i = 0; i < batch_length; ++i) {
            lengths[i] += kByteWidth + kExtraByteForNull;
        }
    }

    void BooleanKeyEncoder::AddLengthNull(int32_t *length) {
        *length += kByteWidth + kExtraByteForNull;
    }

    turbo::Status BooleanKeyEncoder::Encode(const ExecValue &data, int64_t batch_length,
                                            uint8_t **encoded_bytes) {
        if (data.is_array()) {
            VisitArraySpanInline<BooleanType>(
                    data.array,
                    [&](bool value) {
                        auto &encoded_ptr = *encoded_bytes++;
                        *encoded_ptr++ = kValidByte;
                        *encoded_ptr++ = value;
                    },
                    [&] {
                        auto &encoded_ptr = *encoded_bytes++;
                        *encoded_ptr++ = kNullByte;
                        *encoded_ptr++ = 0;
                    });
        } else {
            const auto &scalar = data.scalar_as<BooleanScalar>();
            bool value = scalar.is_valid && scalar.value;
            for (int64_t i = 0; i < batch_length; i++) {
                auto &encoded_ptr = *encoded_bytes++;
                *encoded_ptr++ = kValidByte;
                *encoded_ptr++ = value;
            }
        }
        return turbo::OkStatus();
    }

    void BooleanKeyEncoder::EncodeNull(uint8_t **encoded_bytes) {
        auto &encoded_ptr = *encoded_bytes;
        *encoded_ptr++ = kNullByte;
        *encoded_ptr++ = 0;
    }

    turbo::Result<std::shared_ptr<ArrayData>> BooleanKeyEncoder::Decode(uint8_t **encoded_bytes,
                                                                        int32_t length,
                                                                        MemoryPool *pool) {
        std::shared_ptr<Buffer> null_buf;
        int32_t null_count;
        TURBO_RETURN_NOT_OK(DecodeNulls(pool, length, encoded_bytes, &null_buf, &null_count));

        TURBO_MOVE_OR_RAISE(auto key_buf, allocate_bitmap(length, pool));

        uint8_t *raw_output = key_buf->mutable_data();
        memset(raw_output, 0, bit_util::BytesForBits(length));
        for (int32_t i = 0; i < length; ++i) {
            auto &encoded_ptr = encoded_bytes[i];
            bit_util::SetBitTo(raw_output, i, encoded_ptr[0] != 0);
            encoded_ptr += 1;
        }

        return ArrayData::create(boolean(), length, {std::move(null_buf), std::move(key_buf)},
                                 null_count);
    }

    void FixedWidthKeyEncoder::AddLength(const ExecValue &, int64_t batch_length,
                                         int32_t *lengths) {
        for (int64_t i = 0; i < batch_length; ++i) {
            lengths[i] += byte_width_ + kExtraByteForNull;
        }
    }

    void FixedWidthKeyEncoder::AddLengthNull(int32_t *length) {
        *length += byte_width_ + kExtraByteForNull;
    }

    turbo::Status FixedWidthKeyEncoder::Encode(const ExecValue &data, int64_t batch_length,
                                               uint8_t **encoded_bytes) {
        if (data.is_array()) {
            ArraySpan viewed = data.array;
            auto view_ty = fixed_size_binary(byte_width_);
            viewed.type = view_ty.get();
            VisitArraySpanInline<FixedSizeBinaryType>(
                    viewed,
                    [&](std::string_view bytes) {
                        auto &encoded_ptr = *encoded_bytes++;
                        *encoded_ptr++ = kValidByte;
                        memcpy(encoded_ptr, bytes.data(), byte_width_);
                        encoded_ptr += byte_width_;
                    },
                    [&] {
                        auto &encoded_ptr = *encoded_bytes++;
                        *encoded_ptr++ = kNullByte;
                        memset(encoded_ptr, 0, byte_width_);
                        encoded_ptr += byte_width_;
                    });
        } else {
            const auto &scalar = data.scalar_as<nebula::internal::PrimitiveScalarBase>();
            if (scalar.is_valid) {
                const std::string_view data = scalar.view();
                DKCHECK_EQ(data.size(), static_cast<size_t>(byte_width_));
                for (int64_t i = 0; i < batch_length; i++) {
                    auto &encoded_ptr = *encoded_bytes++;
                    *encoded_ptr++ = kValidByte;
                    memcpy(encoded_ptr, data.data(), data.size());
                    encoded_ptr += byte_width_;
                }
            } else {
                for (int64_t i = 0; i < batch_length; i++) {
                    auto &encoded_ptr = *encoded_bytes++;
                    *encoded_ptr++ = kNullByte;
                    memset(encoded_ptr, 0, byte_width_);
                    encoded_ptr += byte_width_;
                }
            }
        }
        return turbo::OkStatus();
    }

    void FixedWidthKeyEncoder::EncodeNull(uint8_t **encoded_bytes) {
        auto &encoded_ptr = *encoded_bytes;
        *encoded_ptr++ = kNullByte;
        memset(encoded_ptr, 0, byte_width_);
        encoded_ptr += byte_width_;
    }

    turbo::Result<std::shared_ptr<ArrayData>> FixedWidthKeyEncoder::Decode(uint8_t **encoded_bytes,
                                                                           int32_t length,
                                                                           MemoryPool *pool) {
        std::shared_ptr<Buffer> null_buf;
        int32_t null_count;
        TURBO_RETURN_NOT_OK(DecodeNulls(pool, length, encoded_bytes, &null_buf, &null_count));

        TURBO_MOVE_OR_RAISE(auto key_buf, allocate_buffer(length * byte_width_, pool));

        uint8_t *raw_output = key_buf->mutable_data();
        for (int32_t i = 0; i < length; ++i) {
            auto &encoded_ptr = encoded_bytes[i];
            std::memcpy(raw_output, encoded_ptr, byte_width_);
            encoded_ptr += byte_width_;
            raw_output += byte_width_;
        }

        return ArrayData::create(type_, length, {std::move(null_buf), std::move(key_buf)},
                                 null_count);
    }

    turbo::Status DictionaryKeyEncoder::Encode(const ExecValue &data, int64_t batch_length,
                                               uint8_t **encoded_bytes) {
        std::shared_ptr<Array> dict;
        if (data.is_array()) {
            dict = data.array.dictionary().to_array();
        } else {
            dict = data.scalar_as<DictionaryScalar>().value.dictionary;
        }

        if (dictionary_) {
            if (!dictionary_->equals(dict)) {
                // TODO(bkietz) unify if necessary. For now, just error if any batch's dictionary
                // differs from the first we saw for this key
                return turbo::unimplemented_error("Unifying differing dictionaries");
            }
        } else {
            dictionary_ = std::move(dict);
        }
        if (data.is_array()) {
            return FixedWidthKeyEncoder::Encode(data, batch_length, encoded_bytes);
        } else {
            const std::shared_ptr<Scalar> &index = data.scalar_as<DictionaryScalar>().value.index;
            return FixedWidthKeyEncoder::Encode(ExecValue(index.get()), batch_length,
                                                encoded_bytes);
        }
    }

    turbo::Result<std::shared_ptr<ArrayData>> DictionaryKeyEncoder::Decode(uint8_t **encoded_bytes,
                                                                           int32_t length,
                                                                           MemoryPool *pool) {
        TURBO_MOVE_OR_RAISE(auto data,
                            FixedWidthKeyEncoder::Decode(encoded_bytes, length, pool));

        if (dictionary_) {
            data->dictionary = dictionary_->data();
        } else {
            DKCHECK(type_->id() == Type::DICTIONARY);
            const auto &dict_type = turbo::checked_cast<const DictionaryType &>(*type_);
            TURBO_MOVE_OR_RAISE(auto dict, MakeArrayOfNull(dict_type.get_value_type(), 0));
            data->dictionary = dict->data();
        }

        data->type = type_;
        return data;
    }

    void RowEncoder::init(const std::vector<TypeHolder> &column_types, ExecContext *ctx) {
        ctx_ = ctx;
        encoders_.resize(column_types.size());
        extension_types_.resize(column_types.size());

        for (size_t i = 0; i < column_types.size(); ++i) {
            const bool is_extension = column_types[i].id() == Type::EXTENSION;
            const TypeHolder &type = is_extension
                                     ? turbo::checked_pointer_cast<ExtensionType>(
                            column_types[i].get_shared_ptr())
                                             ->storage_type()
                                     : column_types[i];

            if (is_extension) {
                extension_types_[i] = turbo::checked_pointer_cast<ExtensionType>(
                        column_types[i].get_shared_ptr());
            }
            if (type.id() == Type::BOOL) {
                encoders_[i] = std::make_shared<BooleanKeyEncoder>();
                continue;
            }

            if (type.id() == Type::DICTIONARY) {
                encoders_[i] =
                        std::make_shared<DictionaryKeyEncoder>(type.get_shared_ptr(), ctx->memory_pool());
                continue;
            }

            if (is_fixed_width(type.id())) {
                encoders_[i] = std::make_shared<FixedWidthKeyEncoder>(type.get_shared_ptr());
                continue;
            }

            if (is_binary_like(type.id())) {
                encoders_[i] =
                        std::make_shared<VarLengthKeyEncoder<BinaryType>>(type.get_shared_ptr());
                continue;
            }

            if (is_large_binary_like(type.id())) {
                encoders_[i] =
                        std::make_shared<VarLengthKeyEncoder<LargeBinaryType>>(type.get_shared_ptr());
                continue;
            }

            // We should not get here
            DKCHECK(false);
        }

        int32_t total_length = 0;
        for (size_t i = 0; i < column_types.size(); ++i) {
            encoders_[i]->AddLengthNull(&total_length);
        }
        encoded_nulls_.resize(total_length);
        uint8_t *buf_ptr = encoded_nulls_.data();
        for (size_t i = 0; i < column_types.size(); ++i) {
            encoders_[i]->EncodeNull(&buf_ptr);
        }
    }

    void RowEncoder::Clear() {
        offsets_.clear();
        bytes_.clear();
    }

    turbo::Status RowEncoder::EncodeAndAppend(const ExecSpan &batch) {
        if (offsets_.empty()) {
            offsets_.resize(1);
            offsets_[0] = 0;
        }
        size_t length_before = offsets_.size() - 1;
        offsets_.resize(length_before + batch.length + 1);
        for (int64_t i = 0; i < batch.length; ++i) {
            offsets_[length_before + 1 + i] = 0;
        }

        for (int i = 0; i < batch.num_values(); ++i) {
            encoders_[i]->AddLength(batch[i], batch.length, offsets_.data() + length_before + 1);
        }

        int32_t total_length = offsets_[length_before];
        for (int64_t i = 0; i < batch.length; ++i) {
            total_length += offsets_[length_before + 1 + i];
            offsets_[length_before + 1 + i] = total_length;
        }

        bytes_.resize(total_length);
        std::vector<uint8_t *> buf_ptrs(batch.length);
        for (int64_t i = 0; i < batch.length; ++i) {
            buf_ptrs[i] = bytes_.data() + offsets_[length_before + i];
        }

        for (int i = 0; i < batch.num_values(); ++i) {
            TURBO_RETURN_NOT_OK(encoders_[i]->Encode(batch[i], batch.length, buf_ptrs.data()));
        }

        return turbo::OkStatus();
    }

    turbo::Result<ExecBatch> RowEncoder::Decode(int64_t num_rows, const int32_t *row_ids) {
        ExecBatch out({}, num_rows);

        std::vector<uint8_t *> buf_ptrs(num_rows);
        for (int64_t i = 0; i < num_rows; ++i) {
            buf_ptrs[i] = (row_ids[i] == kRowIdForNulls()) ? encoded_nulls_.data()
                                                           : bytes_.data() + offsets_[row_ids[i]];
        }

        out.values.resize(encoders_.size());
        for (size_t i = 0; i < encoders_.size(); ++i) {
            TURBO_MOVE_OR_RAISE(
                    auto column_array_data,
                    encoders_[i]->Decode(buf_ptrs.data(), static_cast<int32_t>(num_rows),
                                         ctx_->memory_pool()));

            if (extension_types_[i] != nullptr) {
                TURBO_MOVE_OR_RAISE(out.values[i], ::nebula::internal::get_array_view(
                        column_array_data, extension_types_[i]));
            } else {
                out.values[i] = column_array_data;
            }
        }

        return out;
    }

}  // namespace nebula::compute::internal
